#include <ros/ros.h>
#include <sensor_msgs/JointState.h>
#include <geometry_msgs/Vector3.h>

#include "eigen3/Eigen/Eigen"
#include "mavlink/v2.0/common/mavlink.h"

#include <cstring>
#include <arpa/inet.h>
#include <unistd.h>

#include <mutex>

#include <cmath>

#define LOCAL_PORT  20000
#define TARGET_IP   "192.168.1.246"
#define TARGET_PORT 14550
#define BUFFER_SIZE 2048

#define M_PI        3.14159265358979323846

typedef struct {
    float w, x, y, z;
} Quaternion_t;

typedef struct {
    float roll, pitch, yaw;
} EulerAngles;

class GimbalController
{
public:
    GimbalController(ros::NodeHandle &nh) : nh_(nh)
    {
        // 创建 UDP 套接字
        sockfd_ = socket(AF_INET, SOCK_DGRAM, 0);
        if (sockfd_ < 0) {
            ROS_ERROR("Error creating socket");
            return;
        }

        // 设置本地地址
        memset(&local_addr_, 0, sizeof(local_addr_));
        local_addr_.sin_family      = AF_INET;
        local_addr_.sin_port        = htons(LOCAL_PORT);
        local_addr_.sin_addr.s_addr = INADDR_ANY;
        if (bind(sockfd_, (struct sockaddr *)&local_addr_, sizeof(local_addr_)) < 0) {
            ROS_ERROR("Error binding socket");
            close(sockfd_);
            return;
        }

        // 设置目标地址
        memset(&target_addr_, 0, sizeof(target_addr_));
        target_addr_.sin_family      = AF_INET;
        target_addr_.sin_port        = htons(TARGET_PORT);
        target_addr_.sin_addr.s_addr = inet_addr(TARGET_IP);

        gimbal_v_[0] = 0.0f;
        gimbal_v_[1] = 0.0f;
        init_pitch = M_PI / 2.0;
        init_yaw = M_PI;
        gimbal_p_[0] = init_pitch;
        gimbal_p_[1] = init_yaw;

        ctrl_mode = "pos_ctrl";

        // 定时器
        recv_timer_ = nh_.createTimer(ros::Duration(0.001), &GimbalController::receive_data, this);
        send_timer_ = nh_.createTimer(ros::Duration(0.005), &GimbalController::send_data, this);


        // 订阅控制命令
        cmd_sub_ = nh_.subscribe("/gimbal/ctrl_cmd", 10, &GimbalController::CtrlCmdCallback, this);
        // 来自IGLOV的控制命令
        cmd1_sub_ = nh_.subscribe("best_view_robot", 10, &GimbalController::iglovCallback, this);

        // 发布关节状态
        joint_state_pub_ = nh_.advertise<sensor_msgs::JointState>("/gimbal/joint_states", 10);
    }

    ~GimbalController()
    {
        close(sockfd_);
    }

private:
    void receive_data(const ros::TimerEvent &)
    {
        char buffer[BUFFER_SIZE];
        int n = recvfrom(sockfd_, buffer, BUFFER_SIZE, 0, nullptr, nullptr);
        if (n < 0) {
            ROS_ERROR("Error receiving data");
            return;
        }
        mavlink_message_t msg;
        mavlink_status_t status;
        for (int i = 0; i < n; ++i) {
            if (mavlink_parse_char(MAVLINK_COMM_0, buffer[i], &msg, &status)) {
                if (msg.msgid == MAVLINK_MSG_ID_HEARTBEAT) {
                    ROS_INFO("Received HEARTBEAT message!!!!");
                } else if (msg.msgid == MAVLINK_MSG_ID_GIMBAL_DEVICE_ATTITUDE_STATUS) {
                    mavlink_gimbal_device_attitude_status_t gimbal_status;
                    mavlink_msg_gimbal_device_attitude_status_decode(&msg, &gimbal_status);

                    sensor_msgs::JointState joint_state_msg;
                    joint_state_msg.header.stamp = ros::Time::now();
                    joint_state_msg.name.push_back("pitch");
                    joint_state_msg.name.push_back("yaw");

                    Quaternion_t gimbal_q;
                    gimbal_q.w = gimbal_status.q[0];
                    gimbal_q.x = gimbal_status.q[1];
                    gimbal_q.y = gimbal_status.q[2];
                    gimbal_q.z = gimbal_status.q[3];
                    auto gimbal_euler = QuaternionToEuler(gimbal_q);

                    joint_state_msg.position.push_back(gimbal_euler.pitch + M_PI / 2.0);
                    joint_state_msg.position.push_back(gimbal_euler.yaw + M_PI);

                    joint_state_msg.velocity.push_back(gimbal_status.angular_velocity_y);
                    joint_state_msg.velocity.push_back(gimbal_status.angular_velocity_z);

                    joint_state_pub_.publish(joint_state_msg);
                }
            }
        }
    }

    void send_data(const ros::TimerEvent &)
    {
        mavlink_message_t msg;
        mavlink_gimbal_device_set_attitude_t attitude_msg;
        attitude_msg.flags              = 0;
        attitude_msg.angular_velocity_x = 0.0f;
        attitude_msg.angular_velocity_y = gimbal_v_[0];
        attitude_msg.angular_velocity_z = gimbal_v_[1];

        Eigen::Quaterniond ctrl_q = Eigen::AngleAxisd(gimbal_p_[1] - M_PI, Eigen::Vector3d::UnitZ()) *
                                    Eigen::AngleAxisd(gimbal_p_[0] - M_PI / 2.0, Eigen::Vector3d::UnitY()) *
                                    Eigen::AngleAxisd(0.0, Eigen::Vector3d::UnitX());
        if (ctrl_mode == "pos_ctrl") {
            attitude_msg.q[0] = ctrl_q.w();
            attitude_msg.q[1] = ctrl_q.x();
            attitude_msg.q[2] = ctrl_q.y();
            attitude_msg.q[3] = ctrl_q.z();
        } else if (ctrl_mode == "vel_ctrl") {
            attitude_msg.q[0] = 0.0f;
            attitude_msg.q[1] = 0.0f;
            attitude_msg.q[2] = 0.0f;
            attitude_msg.q[3] = 0.0f;
        }
        
        attitude_msg.target_system    = 1;
        attitude_msg.target_component = 154;

        mavlink_msg_gimbal_device_set_attitude_encode(1, 154, &msg, &attitude_msg);
        uint8_t buf[MAVLINK_MAX_PACKET_LEN];
        uint16_t len = mavlink_msg_to_send_buffer(buf, &msg);
        if (sendto(sockfd_, buf, len, 0, (struct sockaddr *)&target_addr_, sizeof(target_addr_)) < 0) {
            ROS_ERROR("Sendto failed.");
        }
    }

    // pitch的范围在（1，3.14），yaw的范围在（0，2*M_PI）不然会出错
    void adjust_pitch_yaw(double & pitch, double & yaw){
        pitch = pitch < 1.0 ? 1.0 : pitch;
        pitch = pitch > 3.14 ? 3.14 : pitch;
        while(yaw > 2*M_PI){
            yaw -= 2*M_PI;
        }
        while(yaw < 0){
            yaw += 2*M_PI;
        }
    }

    
    void iglovCallback(const geometry_msgs::Vector3::ConstPtr &msg)
    {
        std::lock_guard<std::mutex> lock(gimbal_mtx_);
        // 程序运行时会设定云台位置为init_pitch何init_yaw，但是iglov返回的角度是在初始状态下计算的，所以这里要加上设定的init角度
        double pitch = msg->y + init_pitch;
        pitch = 0;
        double yaw = -1.0 * msg->z + init_yaw;
        if (std::isnan(yaw) || std::isnan(pitch)) {
            return;
        }
        ctrl_mode    = "pos_ctrl";
        gimbal_p_[0] = pitch;
        gimbal_p_[1] = yaw;
        adjust_pitch_yaw(gimbal_p_[0], gimbal_p_[1]);
        gimbal_v_[0] = 0.1f;
        gimbal_v_[1] = 0.1f;
        ROS_INFO("pos_ctrl!");
    }

    void CtrlCmdCallback(const sensor_msgs::JointState::ConstPtr &msg)
    {
        std::lock_guard<std::mutex> lock(gimbal_mtx_);
        if (msg->name.size() != 2) {
            ROS_ERROR("Wrong ctrl cmd!");
            return;
        }

        
        if (msg->name[0] == "pitch" && msg->name[1] == "yaw") {
            if (msg->position.size() == 2) {
                ctrl_mode    = "pos_ctrl";
                gimbal_p_[0] = msg->position[0];
                gimbal_p_[1] = msg->position[1];
                adjust_pitch_yaw(gimbal_p_[0], gimbal_p_[1]);
                gimbal_v_[0] = msg->velocity.size() == 2 ? msg->velocity[0] : 0.0f;
                gimbal_v_[1] = msg->velocity.size() == 2 ? msg->velocity[1] : 0.0f;
                ROS_INFO("pos_ctrl!");
            } else if (msg->position.empty() && msg->velocity.size() == 2) {
                ctrl_mode    = "vel_ctrl";
                gimbal_p_[0] = M_PI / 2.0;
                gimbal_p_[1] = M_PI;
                gimbal_v_[0] = msg->velocity[0];
                gimbal_v_[1] = msg->velocity[1];
                ROS_INFO("vel_ctrl!");
            } else {
                ROS_ERROR("Wrong ctrl cmd!");
            }
        } else {
            ROS_ERROR("Wrong ctrl cmd!");
        }
    }

    EulerAngles QuaternionToEuler(Quaternion_t q)
    {
        EulerAngles angles;
        float sinr_cosp = 2 * (q.w * q.x + q.y * q.z);
        float cosr_cosp = 1 - 2 * (q.x * q.x + q.y * q.y);
        angles.roll     = atan2(sinr_cosp, cosr_cosp);

        float sinp = 2 * (q.w * q.y - q.z * q.x);
        angles.pitch = fabs(sinp) >= 1 ? copysign(M_PI / 2, sinp) : asin(sinp);

        float siny_cosp = 2 * (q.w * q.z + q.x * q.y);
        float cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z);
        angles.yaw      = atan2(siny_cosp, cosy_cosp);

        return angles;
    }

    ros::NodeHandle nh_;

    int sockfd_;
    struct sockaddr_in local_addr_;
    struct sockaddr_in target_addr_;

    std::string ctrl_mode;
    double gimbal_v_[2];
    double gimbal_p_[2];
    double init_pitch, init_yaw;

    ros::Timer recv_timer_, send_timer_;
    ros::Publisher joint_state_pub_;
    ros::Subscriber cmd_sub_;
    ros::Subscriber cmd1_sub_;

    std::mutex gimbal_mtx_;
};

int main(int argc, char **argv)
{
    ros::init(argc, argv, "nbman_gimbal_controller");
    ros::NodeHandle nh;

    GimbalController controller(nh);

    ros::spin();

    return 0;
}